package com.tianji.aigc.memory;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.tianji.aigc.entity.RedisMessage;
import com.tianji.aigc.result.ToolResultHolder;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.time.temporal.ValueRange;
import java.util.List;
import java.util.Map;

@Component
public class RedisChatMemory implements ChatMemory {

    @Autowired
    private StringRedisTemplate redisTemplate;
    @Override
    public void add(String conversationId, List<Message> messages) {
        //获取redisID
        String key = this.getKey(conversationId);
        //将message进行序列化
        List<String> list = messages.stream().map(message -> {
            RedisMessage redisMessage = BeanUtil.toBean(message, RedisMessage.class);
            redisMessage.setTextContent(message.getText());
            if (message instanceof AssistantMessage assistantMessage){
                String messageId = message.getMetadata().get("id").toString();
                String requestId = (String) ToolResultHolder.get(messageId, "request_id");
                if (StrUtil.isNotBlank(requestId)) {
                    Map<String, Object> paramResult = ToolResultHolder.get(requestId);
                    redisMessage.setParams(paramResult);

                }

                redisMessage.setToolCalls(assistantMessage.getToolCalls());
            }else if(message instanceof ToolResponseMessage  toolResponseMessage){
                redisMessage.setToolResponses(toolResponseMessage.getResponses());

            }
            return JSONUtil.toJsonStr(redisMessage);

        }).toList();

        //将数据存入redis中
        redisTemplate.opsForList().leftPushAll(key, list);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        String key = this.getKey(conversationId);
        //从Redis的List中获取消息字符串列表
        List<String> messageStrList = redisTemplate.opsForList().range(key, 0, lastN);
        if(CollectionUtil.isEmpty(messageStrList)){
            return List.of();
        }
        List<Message> messageList = messageStrList.stream().map(m -> {

            RedisMessage redisMessage = JSONUtil.toBean(m, RedisMessage.class);
            Message message = null;
            //根据消息类型封装具体消息对象
            if(redisMessage.getMessageType().equals(MessageType.USER)){
                message=new UserMessage(redisMessage.getMessageType(),redisMessage.getTextContent(),redisMessage.getMedia(),redisMessage.getMetadata());
            } else if (redisMessage.getMessageType().equals(MessageType.ASSISTANT)){
                message=new AssistantMessage(redisMessage.getTextContent(),redisMessage.getMetadata(),redisMessage.getToolCalls(),redisMessage.getMedia());
            } else if (redisMessage.getMessageType().equals(MessageType.SYSTEM) ){
                message=new SystemMessage(redisMessage.getTextContent());
            } else if (redisMessage.getMessageType().equals(MessageType.TOOL)) {
                message=new ToolResponseMessage(redisMessage.getToolResponses(),redisMessage.getMetadata());
            }
            return message;
        }).toList();
        return messageList;
    }

    @Override
    public void clear(String conversationId) {
        String key = this.getKey(conversationId);
        redisTemplate.unlink(key);

    }


    private String getKey(String conversationId) {
        return "CHAT:" + conversationId;
    }
}
